# add run-length column to exist csv files
import os
import csv
import ast


def run_length_encoding(sequence):
    if not sequence:
        return []

    encoded_sequence = []
    count = 1

    for i in range(1, len(sequence)):
        if sequence[i] == sequence[i - 1]:
            count += 1
        else:
            encoded_sequence.append(count)
            count = 1

    encoded_sequence.append(count)
    if sequence[0] != 1:
        encoded_sequence.insert(0, 0)

    return encoded_sequence


# Set the directory containing the CSV files
directory = './search-result'
csv_files = [filename for filename in os.listdir(
    directory) if filename.endswith('.csv')]
# Sort the list of CSV files alphabetically
csv_files.sort()

# Loop through each file in the directory
for filename in csv_files:
    new_header = []
    rows_with_new_column = []
    # Open the file and read its contents
    with open(os.path.join(directory, filename), 'r') as csvfile:
        reader = csv.reader(csvfile)
        header = next(reader)
        new_header.extend(header)
        new_header.append("run-length")

        for row in reader:
            if len(row) == 0:  # empty row
                continue
            # get the barker pair
            pair = ast.literal_eval(row[3])
            run_length_1 = run_length_encoding(pair[0])
            run_length_2 = run_length_encoding(pair[1])
            run_length = [run_length_1, run_length_2]
            new_row = row
            new_row.append(run_length)
            rows_with_new_column.append(new_row)

    # Write the updated data back to the CSV file
    with open(os.path.join(directory, filename), 'w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(new_header)
        # Write the rows with the new column
        writer.writerows(rows_with_new_column)
